"""
Author: xubing
Date: 2024-11-21
Description: 基础LLM客户端类
"""

import os
from abc import ABC, abstractmethod
from typing import Any, Dict, Generator, List, Optional, Union

import yaml
from pydantic import BaseModel

from logger.logger_config import setup_logger

# 设置日志记录器
logger = setup_logger(__name__)


class Message(BaseModel):
    role: str
    content: Union[str, List[Dict[str, Any]]]


class ChatResponse(BaseModel):
    content: str
    status: bool = True
    error_msg: Optional[str] = None


def get_model_config(model_name: str) -> Dict[str, Any]:
    """获取模型配置"""
    config_path = os.path.join(os.path.dirname(os.path.dirname(__file__)),
                               "config/model_config.yml")
    with open(config_path, "r", encoding="utf-8") as f:
        config = yaml.safe_load(f)

    if model_name not in config:
        raise ValueError(f"Model {model_name} not found in configuration")

    return config[model_name]


class BaseClient(ABC):
    """LLM客户端基类"""

    def __init__(self, model_name: str):
        self.model_name = model_name
        self.model_conf = get_model_config(model_name)
        self.default_params = self.model_conf.get("default_params", {})

    @abstractmethod
    def chat(self, messages: List[Dict[str, Any]]) -> ChatResponse:
        """同步聊天接口"""
        pass

    @abstractmethod
    def chat_stream(
            self, messages: List[Dict[str,
                                      Any]]) -> Generator[str, None, None]:
        """流式聊天接口"""
        pass


if __name__ == "__main__":
    print("BaseClient module loaded successfully")
    print(
        f"Model config path: {os.path.join(os.path.dirname(os.path.dirname(__file__)), '_config', 'model_config.yml')}"
    )
